#!/usr/bin/python

import sys, re, random

def main():
    global outfile
    outfile = open('output.txt', "w")
    txtFile = open(sys.argv[1])
    inputText = txtFile.read()
    txtFile.close()
    keyFile = open(sys.argv[2])
    key=keyFile.read().replace('\n','')
    keyFile.close()
   
    #step 1 preprocessing
    inputText = re.sub('[^A-Z]+', '', inputText)
    txtSize = len(inputText)
    outfile.write('Preprocessing:\n'+inputText+'\n')

    #step 2 Substitution
    cipherText = []
    i = 0
    for chars in inputText:
        ctValue = ((ord(chars)-65) + (ord(key[i%len(key)])-65))%26
        cipherText.append(chr(ctValue+65))
        i+=1
    outfile.write('\nSubstitution:\n'+''.join(cipherText)+'\n')

    #step 3 padding
    while len(cipherText)%16 != 0:
      cipherText.append('A')
    
    #write padding step to file
    outfile.write('\nPadding:\n')
    dr=0
    ns=0
    for elm in cipherText:
        if dr==4:
            outfile.write('\n')
            dr=0
        if ns==16:
            outfile.write('\n')
            ns=0
        outfile.write(elm)
        dr+=1
        ns+=1
        
    #step 4 shifting rows, building 4x4 blocks and converting elements to hex
    bloks = []
    print_bloks=[]#duplicate blocks for printing steps
    z=0
    while z<len(cipherText):
        cube=[]
        cz=[]
        y=0
        while y<4:
            row=[]
            rw=[]
            x=0
            while x<4:
                row.append(hex(par_bit_flip(ord(cipherText[z])))[2:4])
                rw.append(cipherText[z])
                z+=1
                x+=1
          #shift row
            i=0
            while i<y:
                row.append(row.pop(0))
                rw.append(rw.pop(0))
                i+=1
            cube.append(row)
            cz.append(rw)
            y+=1
        bloks.append(cube)
        print_bloks.append(cz)
        
    #write shift rows to file 
    outfile.write('\n\nShift rows:\n')
    for qa in print_bloks:
        scribble(qa)
    #write parity to file
    outfile.write('\n\nParity:\n')
    for qb in bloks:
        scribble(qb)
    #begin mix columns
    outfile.write('\n\nMix Columns:\n')
    for qb in bloks:
        scribble(mix_columns(qb))
    
    
#step 5: parity bit
#int arg, flips parity bit if odd parity, returns int
def par_bit_flip(val):
    if par_odd(val)==-1:
        val=val^128
    return(val)

#influenced by parityOf "Bit Twiddling Hacks"
#int arg, returns -1 if odd par, 0 if even
def par_odd(i):
    par=0
    while(i):
        par= ~par
        i = i & (i-1)
    return(par)

#step 6: mix columns
def mix_columns(matrix):
    mds=[[2,3,1,1],[1,2,3,1],[1,1,2,3],[3,1,1,2]]
    result=[[0,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]]
    mor=[]
    z=0
    while z<4:
        x=0
        neo=[]
        while x<4:
            y=0
            r=0
            while y<4:
                c=int(matrix[y][z],16)
                r^=gal_mult(mds[x][y], c)
              #  print y,'c:',c,'  mds:',mds[x][y],'r:',r
                y+=1
            result[x][z]=hex(r)[2:4]
            x+=1
        z+=1
    return result

def gal_mult(case, val):
    result = val;
    if case==2:
        result = (127&val)<<1
        if (128|val)==val:
            result^=27
    elif case==3:
        result = (127&val)<<1
        if (128|val)==val:
            result^=27
        result^=val
    return result;

def scribble(blist):
    i=0
    for nlist in blist:
        for c in nlist:
            outfile.write(c+' ')
        outfile.write('\n')
        if i==3:
            outfile.write('\n')
            i=-1
        i+=1
    
if __name__ == '__main__':
    main()
